import Foundation

extension Sequence where Element == Float {
    func mean() -> Float {
        var sum: Float = 0
        var count: Int = 0
        for value in self {
            sum += value
            count += 1
        }
        return count == 0 ? .nan : sum / Float(count)
    }

    func onlyNaN() -> Bool {
        return first(where: { !$0.isNaN }) == nil
    }

    func hasValidData() -> Bool {
        return !onlyNaN()
    }
}

extension Array where Element == Float {
    func max(by: Int) -> [Float] {
        return stride(from: 0, through: count - by, by: by).map { i in
            // Note: `.max()` ignores NaN values
            return self[i..<i + by].reduce(-Float.greatestFiniteMagnitude, {$1 < $0 ? $0 : $1})
        }
    }
    func min(by: Int) -> [Float] {
        return stride(from: 0, through: count - by, by: by).map { i in
            // Note: `.min()` ignores NaN values
            return self[i..<i + by].reduce(Float.greatestFiniteMagnitude, {$1 > $0 ? $0 : $1})
        }
    }
    func sum(by: Int) -> [Float] {
        return stride(from: 0, through: count - by, by: by).map { i in
            return self[i..<i + by].reduce(0, +)
        }
    }
    func sum(by: Float) -> [Float] {
        return stride(from: Float(0), through: Float(count) - by, by: by).map { (i: Float) -> Float in
            return self[Int(roundf(i))..<Int(roundf(i + by))].reduce(0, +)
        }
    }
    func mean(by: Int) -> [Float] {
        return stride(from: 0, through: count - by, by: by).map { i in
            return self[i..<i + by].reduce(0, +) / Float(by)
        }
    }
    func meanBackwards(by: Int) -> [Float] {
        return stride(from: 0, through: count - by, by: by).map { i in
            if i == 0 {
                return .nan
            }
            return self[i - by..<i].reduce(0, +) / Float(by)
        }
    }

    mutating func rounded(digits: Int) {
        let roundExponent = powf(10, Float(digits))
        for i in indices {
            self[i] = Foundation.round(self[i] * roundExponent) / roundExponent
        }
    }

    func round(digits: Int) -> [Float] {
        let roundExponent = powf(10, Float(digits))
        return map {
            return Foundation.round($0 * roundExponent) / roundExponent
        }
    }

    func containsNaN() -> Bool {
        return first(where: { $0.isNaN }) != nil
    }

    mutating func fillWithNaNs() {
        for i in self.indices {
            self[i] = .nan
        }
    }
}

extension Array where Element == Float {
    /// Shift longitudes by 180° and flip south.north
    mutating func shift180LongitudeAndFlipLatitude(nt: Int, ny: Int, nx: Int) {
        precondition(nt * ny * nx == count)
        self.withUnsafeMutableBufferPointer { data in
            /// Data starts eastwards at 0°E... rotate to start at -180°E
            for t in 0..<nt {
                for y in 0..<ny {
                    for x in 0..<nx / 2 {
                        let offset = t * nx * ny + y * nx
                        let val = data[offset + x]
                        data[offset + x] = data[offset + x + nx / 2]
                        data[offset + x + nx / 2] = val
                    }
                }
                /// Also flip south / north
                for y in 0..<ny / 2 {
                    for x in 0..<nx {
                        let val = data[t * nx * ny + y * nx + x]
                        data[t * nx * ny + y * nx + x] = data[t * nx * ny + (ny - 1 - y) * nx + x]
                        data[t * nx * ny + (ny - 1 - y) * nx + x] = val
                    }
                }
            }
        }
    }

    /// Shift longitudes by 180°
    mutating func shift180Longitude(nt: Int, ny: Int, nx: Int) {
        precondition(nt * ny * nx == count)
        self.withUnsafeMutableBufferPointer { data in
            /// Data starts eastwards at 0°E... rotate to start at -180°E
            for t in 0..<nt {
                for y in 0..<ny {
                    for x in 0..<nx / 2 {
                        let offset = t * nx * ny + y * nx
                        let val = data[offset + x]
                        data[offset + x] = data[offset + x + nx / 2]
                        data[offset + x + nx / 2] = val
                    }
                }
            }
        }
    }

    /// Flip south.north
    mutating func flipLatitude(nt: Int, ny: Int, nx: Int) {
        precondition(nt * ny * nx == count)
        self.withUnsafeMutableBufferPointer { data in
            for t in 0..<nt {
                /// Also flip south / north
                for y in 0..<ny / 2 {
                    for x in 0..<nx {
                        let val = data[t * nx * ny + y * nx + x]
                        data[t * nx * ny + y * nx + x] = data[t * nx * ny + (ny - 1 - y) * nx + x]
                        data[t * nx * ny + (ny - 1 - y) * nx + x] = val
                    }
                }
            }
        }
    }

    mutating func multiplyAdd(multiply: Float, add: Float) {
        self.withUnsafeMutableBufferPointer { data in
            for i in 0..<data.count {
                data[i] = fma(data[i], multiply, add)
            }
        }
    }

    /// Perform a simple delta encoding. Leave the first value as seed.
    mutating func deltaEncode() {
        if count <= 1 {
            return
        }
        for x in (1..<count).reversed() {
            if self[x - 1].isNaN {
                continue
            }
            self[x] = self[x - 1] - self[x]
        }
    }

    /// Undo delta coding
    mutating func deltaDecode() {
        if count <= 1 {
            return
        }
        for x in 1..<count {
            if self[x - 1].isNaN {
                continue
            }
            self[x] = self[x - 1] - self[x]
        }
    }

    /// Ensure data is greater equals than
    mutating func greater(than value: Float) {
        for x in 0..<count {
            self[x] = Swift.max(self[x], value)
        }
    }
}
